Added the concrete metadata-contract proposal (the "(1)" item from the PR description): commit 4456a03 adds inference_parameters / context_parameters / standardization as properties on NeuralDistribution — the supported interface for samplers, working on old checkpoints unchanged.

Full proposal for the sampler group (incl. the step-3 dict-batch call-signature migration plan and the exact factors.py before/after): hackathon/NN_Sampler_Interface.md. Group A sign-off on that doc is the remaining coordination point.

Step 3 landed: dict batches, registry-built embeddings, autocomplete_model_kwargs deleted

Three new commits (f917def2, d54a1d33, d2f9315a) complete step 3 of the build-system design (§6). All 321 tests in tests/core + tests/gw pass.

Note on the rebase: the branch was rebased onto the updated hackathon-1. #387 had added a smaller tests/core/test_build_model.py; the conflict was resolved in favor of the more comprehensive version from this branch. One #387 test asserting case-insensitive dispatch was dropped deliberately: per the step-2 decision, registry lookups are case-sensitive, and historical case-tolerance lives only in update_model_config.

1. f917def2 — dict batches (design §4.4, interface doc §2)

• The training transform chain ends in SelectKeys (dict-preserving) instead of UnpackDict (positional flattening); loaders yield {"inference_parameters": ..., "waveform": ..., "context_parameters": ...}.
• NeuralDistribution.sample / sample_and_log_prob / log_prob / loss take a single context: dict | None instead of positional *context — exactly the signature proposed to Group A in NN_Sampler_Interface.md §2.
• FlowWrapper and ContinuousFlow route context[k] for k in context_keys into the embedding network once, at the model boundary; tensors stay positional inside modules. A missing context key raises a ValueError naming it — no more silent dependence on the loader's key order.
• This kills the GNPE-proxy detection via try/except IndexError and the overloaded data_sample[2] slot (which the transformer branch reused for position tensors — the two were mutually exclusive by accident).
• Call sites updated: core/samplers.py (plain + GNPE), SampleDataset for unconditional training.

2. d54a1d33 — new schema + registries; autocomplete_model_kwargs deleted (design §4.3/4.5/4.7/4.8)

New model-settings schema, each component a registry type plus a flat kwargs dict:

model:
  distribution:    {type: normalizing_flow, kwargs: {num_flow_steps: 30, ...}}
  embedding_net:   {type: dense_svd, kwargs: {output_dim: 128, ...}}
  # context_merger: {type: concat}   # added automatically when the data has context_parameters

• Embedding contract (core/nn/enets.py module docstring): input_keys, output_dim, complete_settings(settings, sample_batch). First registered components: dense_svd (LinearProjectionRB + DenseResidualNet) and the concat context merger. Any embedding now composes with any distribution type — transformer×FMPE will need zero extra wiring.
• complete_model_settings replaces autocomplete_model_kwargs: each architecture infers its own input dims from the named sample batch; generic theta_dim/context_dim are computed once; completed settings are saved in the checkpoint so loading never needs a data sample. Dimensions in user YAML are an error; a context_merger with no context parameters in the data is an error (no silent drops).
• initialize_network loses its embedding_kwargs-presence branches; NeuralDistribution.build_embedding_net assembles embedding + merger generically. The V_rb_list deepcopy hack is gone — initial weights are extra constructor kwargs, never in saved settings.
• save_model enforces the parameter contract promised to samplers (interface doc §1): standardization must cover inference ∪ context parameters, for built-ins and plugins alike.
• Back-compat: update_model_config is the single boundary and maps all old schemas forward (nsf+embedding, posterior_model_type/posterior_kwargs/embedding_kwargs, added_context → concat merger, input_dim → theta_dim). Old checkpoints build state-dict-identical networks — pinned by test_old_schema_builds_state_dict_compatible_network and an actual old-format .pt load round trip (test_load_old_schema_checkpoint_file).
• Deleted: the never-wired embedding_net_builder parameter, create_nsf_wrapped, create_nsf_with_rb_projection_embedding_net. Example configs, pipe defaults, and the architecture docs notebook use the new schema (old YAML still works through the shim).

3. d2f9315a — chained-NPE training-side conditioning (design §4.5)

• data.conditioning_parameters feed the same context_parameters list as GNPE proxies — two sources, one declared contract.
• ContextMergerMLP ported as the registered mlp merger: context mixed in through a learned MLP so context_dim doesn't grow; its dims are inferred at completion (no hand-written input_dim), output_dim defaults to the embedding's.
• data.time_alignment: true trains on time-aligned data via ProjectOntoDetectors(apply_time_shift=False), with the validation guard (requires ra/dec/geocent_time as conditioning, forbids them as inference targets, incompatible with gnpe_time_shifts). Tests ported from the chained branch. The chained inference side is a factorized-sampler chain (Group A's domain) and is not built here.

Deliberately deferred

• FlowFactor/GNPEFlowFactor call sites (interface doc §2, impacts 1–3): factors.py is not in hackathon-1 yet, so those edits can't land in this PR. Nothing on sampler-revamp breaks in the meantime; we'll do that migration as promised once the branches meet.
• Step 4 (architecture-owned SVD init — the initial_weights threading and the trainer's svd gate still exist, marked with a TODO) and steps 5–6 (transformer merge) follow.

Review pointers

• The schema mapping: core/utils/backward_compatibility.py::update_model_config
• The autocomplete replacement: core/posterior_models/build_model.py::complete_model_settings
• The embedding contract + registered components: core/nn/enets.py
• Checkpoint compatibility tests: tests/core/test_build_model.py (old-schema section)

🤖 Generated with Claude Code

Steps 4–6 landed: init hook, transformer merge, T1 augmentation — the plan is complete

Three further commits (03838a12, 08d9dbe9, 87888a59) finish the build-system design §6 sequence. All 409 tests in tests/core + tests/gw pass. The PR now carries the full refactor; implementation deltas from the design sketches are footnoted in the design doc (§4.3, §4.6).

4. 03838a12 — architecture-owned initialization (design §4.6)

• Two optional hooks on the embedding contract: init_data_spec() declares the data variation a network wants for weight initialization (spec keys: noise, network_format, fix_parameters, num_samples), and initialize_weights(batches, out_dir) consumes a matching batch iterator. DenseSVDEmbedding implements them, absorbing all of build_svd_for_embedding_network (incl. validation diagnostics; V-matrix truncation now uses the network's own num_bins, so core stays domain-agnostic).
• The trainer is generic: prepare_training_new walks pm.network.modules() and answers any module's spec via initialization_dataloader (train_builders), which maps the spec onto the existing omit_transforms mechanism. Hooks compose through wrappers (mergers, FlowWrapper) and plugins get data-driven init without touching dingo. Seeding triggers on svd.num_training_samples, so loading a saved model never re-runs it.
• Deleted: build_svd_for_embedding_network, the entire initial_weights constructor threading, and the "embedding network is assumed to have an SVD projection layer" trainer special case. Note for the transformer/T1 authors: the initial_weights= constructor argument no longer exists.
• SVDBasis moved dingo/gw/SVD.py → dingo/core/SVD.py (it was always domain-agnostic; the old path re-exports).

5. 08d9dbe9 — transformer merge, stage 1 (from origin/transformer)

• TransformerEmbedding registers as transformer: tokenized strain → position-conditioned Tokenizer → TransformerEncoder → CLS/average pooling → optional final net. complete_settings infers tokenizer dims and num_blocks from the named sample batch — position is a regular batch entry, no longer squatting in the GNPE-proxy slot.
• The seam paid off, pinned by tests: transformer × flow-matching works with zero extra wiring, and transformer + context_parameters composes via the concat merger — both impossible on the original branch.
• DenseResidualNet unified, not duplicated: the branch's block (MyResidualBlock) is a strict superset of glasflow's ResidualBlock — identical parameter names and forward pass at layer_norm=False (pinned by an equivalence test) — so it moved to core/nn/resnet.py and replaced the glasflow-block-based version everywhere. Old checkpoints load unchanged. The one genuinely incompatible case, the rq-coupling conditioner, is selected by an explicit base_transform_kwargs.conditioner_type: glasflow_residual | dense_residual (per the team decision: a distinct type, not a layer_norm flag).
• StrainTokenization wired via data.tokenization in training and GWSampler inference. Not ported: the dead drop_* config keys (→ step 6), the allow_tf32 pop, the embedding_type branches. New examples/transformer_model config.

6. 87888a59 — T1 augmentation + token suppression (from origin/dingo-t1)

• Training augmentation via data.tokenization: DropDetectors, DropFrequenciesToUpdateRange (f_cut), DropFrequencyInterval (mask_interval), DropRandomTokens, NormalizePosition. Two deviations from dingo-t1, both deliberate: DropDetectors gained unbatched-sample support (our pipeline feeds single samples), and DropRandomTokens lost the increase_p_until_epoch ramp — dead on dingo-t1 too (the epoch injection is commented out there).
• Inference: GWSampler.suppress ([f_lo, f_hi] or per-detector dict) with real validation replacing T1's warning stub; tokenized models route frequency updates through the ported UpdateFrequencyRange (token masks instead of zeroed strain). Two T1 quirks fixed: positions are normalized after Hz-threshold comparisons, and unchanged domain bounds are passed as None so the zero-padded final token isn't spuriously masked. UpdateFrequencyRange got the test T1 left as a TODO.
• PoolingTransformer (MultiPositionalEncoding + average pooling) registers as pooling_transformer — a third tokenized architecture through the same generic path.
• Deliberately not ported (separate research features on dingo-t1, need their own decisions): pretraining (enets_pretraining/pretraining_model), scheduler/AMP trainer changes, mixed ASD datasets, TimeShiftStrainGrid, BlockEncoding, skymaps.

Review pointers

• Init hooks + trainer glue: core/nn/enets.py (init_data_spec/initialize_weights), gw/training/train_pipeline.py::prepare_training_new, gw/training/train_builders.py::initialization_dataloader
• Block equivalence (checkpoint compatibility): tests/core/test_resnet.py::test_residual_block_equivalent_to_glasflow
• Transformer contract + cross-distribution tests: core/nn/transformer.py, tests/core/test_build_model.py (transformer section)
• Augmentation transforms + ported tests: gw/transforms/tokenization_transforms.py, tests/gw/transforms/test_tokenization_augmentation.py
• Token suppression: gw/inference/gw_samplers.py (suppress, _initialize_transforms), tests/gw/inference/test_gw_sampler_transforms.py

Open follow-ups (outside this PR)

• factors.py migration to the accessors + dict context once sampler-revamp reaches hackathon-1 (interface doc §2 — we do those edits).
• Hydra structured configs as their own PR (design §9; both constraints honored here).
• T1 leftovers above, to be scoped with the T1 authors.

🤖 Generated with Claude Code